#!/usr/bin/python
# coding=utf-8
"""
SSH to remote server and execute mkdir|upload|download|run|kill|chtime|passwd|rm.
use: python ssh.py mkdir|upload|download|run|kill|chtime|passwd|rm\n\
    mkdir    mkdir /home/OTS-TEST\n\
    upload   upload vmstat.py\n\
    download download result file\n\
    run      start run vmstat.py\n\
    kill     end process vmstat\n\
    chtime   time synchronization\n\
    passwd   change passwd\n\
    rm       remove /home/OTS-TEST\n
    
    The info of remote server should be saved in hosts.txt and the formtat is as follows:
    [ip],[username],[passwd]    
Email: billy7x17@gmail.com
"""

import os
import sys
import time
import paramiko
import threading
import glob

sample_dict = {}
result_dict = {}
key_list = []
rootdir = "/home/OTS-TEST"
default_threshold = 60


def getIP():
    """
    查询hosts*.txt中的主机信息，获取主机dict
        + 通过hosts-组名.txt命名来进行分组
        + 通过hosts-组名(阈值).txt或hosts(阈值).txt命名来设置阈值

    :return: 主机dict
    {分组 : [[ip,user,password], ...], ...}
    """
    data = {}
    read_files = glob.glob("hosts*.txt")
    for i in read_files:
        with open(i, 'r') as infile:
            key_name = infile.name[5:-4]

            # 获取阈值（获取括号中的字符串）
            threshold = default_threshold
            if key_name.endswith(")"):
                threshold = key_name[key_name.index("(") + 1: key_name.index(")")]
                key_name = key_name[:key_name.index("(")]

            # 获取分组
            if key_name.startswith("-"):
                key_name = key_name[1:]
            elif key_name == "":
                key_name = "none"
            line = infile.readline()
            data[key_name] = []
            while line:
                line_arr = line.strip().split(",")
                line_arr.append(threshold)
                data[key_name].append(line_arr)
                key_list.append(key_name + "-" + line_arr[0])
                line = infile.readline()
            infile.close()
    return data

    # f = open("hosts*.txt", "r")
    # line = f.readline()
    # while line:
    #     line = line.strip()
    #     record = line.split(",")
    #     for i in range(0, len(record)):
    #         record[i] = record[i].strip()
    #     key_list.append(record[0])
    #     records.append(record)
    #     line = f.readline()
    # f.close()
    # return records


def con_dict():
    """
    采集所有ip，便于打印失败的主机ip

    :return: none
    """
    result_dict.update(dict.fromkeys(key_list, 0))
    sample_dict.update(dict.fromkeys(key_list, 0))


def update_dict():
    result_dict.update(sample_dict)


def check_result():
    """
    打印执行命令失败的主机
    :return: none
    """
    cnt = 0
    print 'unsuccessful group-ip:'
    for item in result_dict:
        if result_dict[item] == 0:
            cnt += 1
            print item
    if cnt == 0:
        print 'None'
    update_dict()


def ssh2(ip, username, passwd, cmd, group):
    """
    ssh2协议登录主机并执行命令

    :param ip: ip地址
    :param username: 用户名
    :param passwd: 密码
    :param cmd: 执行的命令
    :return:none
    """
    try:
        flag = True
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(ip, 22, username, passwd, timeout=30)
        for m in cmd:
            stdin, stdout, stderr = ssh.exec_command(m)
            if cmd[0].split(' ')[0] != 'passwd':
                err = stderr.readlines()
                out = stdout.readlines()
                if out:
                    for o in out:
                        print o,
                if err:
                    flag = False
                    for e in err:
                        print e,
                else:
                    print 'ip: %s exec cmd <%s> success' % (ip, m)
            time.sleep(1)
            if cmd[0].split(' ')[0] == 'passwd':
                stdin.write('cmri-ots\n')
                time.sleep(2)
                stdin.write('cmri-ots\n')
                time.sleep(2)
                print 'ip: %s exec cmd <%s> success' % (ip, m.strip())
        if not flag:
            raise IOError, IOError(err)
        result_dict[group + "-" + ip] = 1
        # ssh.close()
    except:
        print '%s\tError\n' % (ip)
    finally:
        ssh.close()


def sftp_up(ip, usrname, passwd, mod, group):
    """
    tftp协议登录主机并根据命令上传文件

    :param ip: ip地址
    :param username: 用户名
    :param passwd: 密码
    :param mod: 模式
        - mkdir创建根目录文件夹
        - upload上传采集脚本
    :return:none
    """
    port = 22
    t = None
    try:
        t = paramiko.Transport((ip, port))
        t.connect(username=usrname, password=passwd)
        print 'sftp_up connect success'
        sftp = paramiko.SFTPClient.from_transport(t)
        local_path = '.'
        if mod == 'mkdir':
            remote_path = '/home'
            sftp.put(local_path + '/mkdir.py', remote_path + '/mkdir.py')
            print "ip: %s exec cmd <upload mkdir.py> success" % (ip)
        elif mod == 'upload':
            remote_path = '%s/' % (rootdir)
            sftp.put(local_path + '/vmstat.py', remote_path + '/vmstat.py')
            print "ip: %s exec cmd <upload vmstat.py> success" % (ip)
        result_dict[group + "-" + ip] = 1
        # t.close()
    except:
        import traceback
        traceback.print_exc()
        print "%s error, ip: %s" % (mod, ip)
    finally:
        if t:
            t.close()


def sftp_down(ip, username, passwd, threshold, path):
    """
    sftp协议下载性能结果文件

    :param ip: ip地址
    :param username: 用户名
    :param passwd: 密码
    :param threshold: 阈值
    :param path: 下载到本地的路径
    :return:none
    """
    apath = rootdir
    apattern = '"vmstat*.txt"'
    rawcommand = 'find {path} -name {pattern}'
    command = rawcommand.format(path=apath, pattern=apattern)

    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(ip, 22, username, passwd, timeout=5)
        stdin, stdout, stderr = ssh.exec_command(command)
        filelist = stdout.read().splitlines()

        ftp = ssh.open_sftp()
        if not os.path.exists(path):
            os.mkdir(path, 0700)
        thresholdFile = open(path + "/.sshvmstat", 'w+')
        thresholdFile.write("threshold:%s" % threshold)
        for afile in filelist:
            (head, filename) = os.path.split(afile)
            ftp.get(afile, path + '/' + filename)
        print "ip: %s exec cmd <download vmstat_xxx.txt> success" % (ip)
        result_dict[path + "-" + ip] = 1
        # ftp.close()
        # ssh.close()
    except:
        import traceback
        traceback.print_exc()
        print "download error, ip: %s" % (ip)
    finally:
        ftp.close()
        ssh.close()


def passwd(ip, usrname, passwd):
    port = 22
    try:
        t = paramiko.Transport((ip, port))
        t.connect(username=usrname, password=passwd)
        print 'passwd connect success'
        s = t.open_session()
        s.get_pty()
        s.invoke_shell()
        s.send('passwd root\n')
        time.sleep(2)
        s.send('123456\n')
        time.sleep(2)
        s.send('123456\n')
        time.sleep(2)
        result_dict[ip] = 1
    except Exception, err:
        print "Err:unable to connect %s:%s\n" % (ip, err)


def exec_cmd(records, mod):
    """
    批量执行命令

    :param records: 主机dict
    :param mod: 模式
    :return: none
    """
    cmd = []
    if mod == 'mkdir':
        cmd_exec = 'python /home/mkdir.py'
        cmd.append(cmd_exec)
        cmd_exec = 'rm -rf /home/mkdir.py'
    elif mod == 'run':
        # cmd_exec = "nohup python %s/vmstat.py >/dev/null 2>&1 &" %(rootdir)
        cmd_exec = "nohup python %s/vmstat.py >%s/nohup.txt 2>&1 &" % (rootdir, rootdir)
    elif mod == 'kill':
        cmd_exec = "ps -ef | grep vmstat | grep -v grep | cut -c 9-15| xargs kill -9"
    elif mod == 'chtime':
        now_time = os.popen('date "+%Y-%m-%d %H:%M:%S"').readlines()
        cmd_exec = 'date -s ' + '"' + now_time[0].strip() + '"'
    elif mod == 'passwd':
        cmd_exec = 'passwd root\n'
    elif mod == 'rm':
        cmd_exec = 'rm -rf %s' % (rootdir)
    cmd.append(cmd_exec)
    for fileKey in records.iterkeys():
        threads = []
        lineKey = records[fileKey]
        print "%s:" % fileKey
        for r in lineKey:
            a = threading.Thread(target=ssh2, args=(r[0], r[1], r[2], cmd, fileKey))
            threads.append(a)
        for a in threads:
            a.setDaemon(True)
            a.start()
        for a in threads:
            a.join()


def trans_file(records, mod):
    """
    文件传输

    :param records: 主机dict
    :param mod: 模式
    :return: none
    """
    for fileKey in records.iterkeys():
        threads = []
        lineKey = records[fileKey]
        print "%s:" % fileKey
        for r in lineKey:
            if mod == 'upload' or mod == 'mkdir':
                a = threading.Thread(target=sftp_up, args=(r[0], r[1], r[2], mod, fileKey))
            elif mod == 'download':
                a = threading.Thread(target=sftp_down, args=(r[0], r[1], r[2], r[3], fileKey))
            elif mod == 'passwd':
                a = threading.Thread(target=passwd, args=(r[0], r[1], r[2]))
            threads.append(a)

        for a in threads:
            a.setDaemon(True)
            a.start()
        for a in threads:
            a.join()


# main 方法
if __name__ == '__main__':
    records = getIP()
    con_dict()
    action = ['upload', 'download', 'run', 'kill', 'mkdir', 'chtime', 'passwd', 'rm']
    if (len(sys.argv) != 2) or (sys.argv[1] not in action):
        print 'use: python ssh.py mkdir|upload|download|run|kill|chtime|passwd|rm\n\
            mkdir    mkdir /home/OTS-TEST\n\
            upload   upload vmstat.py\n\
            download download result file\n\
            run      start run vmstat.py\n\
            kill     end process vmstat\n\
            chtime   time synchronization\n\
            passwd   change passwd\n\
            rm       remove /home/OTS-TEST\n'
    elif sys.argv[1] == 'upload':
        trans_file(records, 'upload')
    elif sys.argv[1] == 'download':
        trans_file(records, 'download')
    elif sys.argv[1] == 'run':
        exec_cmd(records, 'run')
    elif sys.argv[1] == 'kill':
        exec_cmd(records, 'kill')
    elif sys.argv[1] == 'mkdir':
        trans_file(records, 'mkdir')
        print "upload mkdir.py success"
        exec_cmd(records, 'mkdir')
    elif sys.argv[1] == 'chtime':
        exec_cmd(records, 'chtime')
    elif sys.argv[1] == 'passwd':
        # trans_file(records, 'passwd')
        exec_cmd(records, 'passwd')
    elif sys.argv[1] == 'rm':
        exec_cmd(records, 'rm')
    check_result()
